import pandas as pd
import os
import random  # 导入random模块


def extract_sample(input_path, output_path=None, sample_fraction=0.1, random_state=42):
    """
    从CSV文件中提取指定比例的样本数据并保存到新文件

    参数:
    input_path: 原始CSV文件路径
    output_path: 输出文件路径，默认为原文件名加"_sample"后缀
    sample_fraction: 抽样比例，默认为0.1(10%)
    random_state: 随机种子，保证结果可重现
    """
    # 设置随机种子，确保抽样结果可重现
    random.seed(random_state)

    # 确定输出文件路径
    if output_path is None:
        dir_name, file_name = os.path.split(input_path)
        base_name, ext = os.path.splitext(file_name)
        output_path = os.path.join(dir_name, f"{base_name}1{ext}")

    try:
        # 读取CSV文件
        print(f"正在读取文件: {input_path}")

        # 对于大型CSV，先获取总行数，再计算需要抽样的行数
        with open(input_path, 'r', encoding='utf-8') as f:
            total_rows = sum(1 for _ in f)

        # 计算需要抽样的行数（保留整数）
        sample_size = int(total_rows * sample_fraction)
        print(f"总数据行数: {total_rows}")
        print(f"将提取 {sample_size} 行数据 ({sample_fraction * 100}%)")

        # 随机抽样并保存
        skip = sorted(random.sample(range(1, total_rows), total_rows - sample_size))
        df = pd.read_csv(input_path, skiprows=skip, low_memory=False)

        # 保存抽样数据
        df.to_csv(output_path, index=False)
        print(f"抽样数据已保存至: {output_path}")

        return output_path

    except Exception as e:
        print(f"处理过程中出错: {str(e)}")
        return None


if __name__ == "__main__":
    # 示例用法
    input_csv = r"C:\Users\LGXY\Downloads\UserBehavior.csv"  # 替换为你的原始文件路径
    extract_sample(input_csv)
